#include "deepsort/SORT.h"

using namespace std;

struct SORT::TrackData {
    KalmanTracker kalman;
};

SORT::SORT()
    : manager(make_shared<TrackerManager<TrackData>>(data)),
    drawer(make_shared<MotDrawer>()){}

SORT::~SORT() = default;

Result SORT::Process(cv::Mat& ori_img, std::vector<BBox>& detections, std::vector<BBox>& trackResults) {
    manager->predict();
    manager->remove_nan();

    auto metric = [this, &detections](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) {
        vector<BBox> trks;
        for (auto t : trk_ids) {
            BBox b;
            b.rect = data[t].kalman.rect();
            trks.push_back(b);
        }
        vector<BBox> dets;
        for (auto &d:det_ids) {
            dets.push_back(detections[d]);
        }

        auto iou_mat = iou_dist(dets, trks);
        for (size_t i = 0; i < iou_mat.size(); ++i) {
            for (size_t j = 0; j < iou_mat[0].size(); ++j) {
                if (iou_mat[i][j] > 0.7f) {
                    iou_mat[i][j] = INVALID_DIST;
                }
            }
        }
        return iou_mat;
    };
    auto matched = manager->update(detections, metric, metric);
    manager->remove_deleted();

    manager->visible_tracks(trackResults);


    // Utils::WriteImageToFile(ori_img, "output_sort");

    return SUCCESS;
}

void SORT::DrawTrajectories(cv::Mat& ori_img, vector<BBox>& trackResults, bool isDrawTrajectory) 
{
    drawer->DrawTrajectories(ori_img, trackResults, isDrawTrajectory);
}

Result SORT::Set(uint32_t width, uint32_t height) {
    manager->Set(width, height);
    return SUCCESS;
}
